# 8个线性天线一层2次多项式拟合：12天线、sslDb=-30dB、扫描角度0度
import argparse
import sys
from typing import Dict
import numpy as np
from scipy.signal import windows as SSW # import chebwin
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.optim.lr_scheduler import ReduceLROnPlateau
from apps.wfs.nlbf_config import NlbfConfig as NG

class PolynomialModel1(nn.Module):
    def __init__(self, N_base:int =8, rank:int =3, init_coefs:bool =False, avec:torch.tensor =None, avec_dim:int =1):
        super(PolynomialModel1, self).__init__()
        self.rank = rank
        self.N_base = N_base
        for idx in range(avec.shape[0]):
            # a 的形状是 (8,)
            self.a_powers = torch.stack([avec[idx]**i for i in range(self.rank)], dim=1)
        # 定义一个8x4的矩阵，每一行对应一个多项式的系数
        if init_coefs:
            self.load()
        else:
            for i in range(avec_dim):
                self.register_parameter(f'c_{i}', nn.Parameter(torch.randn(N_base, self.rank)))
    
    def forward(self, a):
        '''
        a: (1800, 8)
        '''
        outputs = None
        for idx in range(a.shape[0]):
            # a 的形状是 (8,)
            a_power = torch.stack([a[idx]**i for i in range(self.rank)], dim=1)
            # a_power 的形状是 (8, 4)
            # 计算每个多项式的值
            polynomial_values = torch.sum(getattr(self, f'c_{0}') * a_power, dim=1)
            # polynomial_values 的形状是 (8,)
            # 求和
            output = torch.abs(torch.sum(polynomial_values))
            if outputs is None:
                outputs = output
            else:
                outputs = torch.vstack((outputs, output))
        return outputs

    def save(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        torch.save(self.state_dict(), pt_fn)

    def load(self, pt_fn='./work/wfs/nlbf.pt') -> None:
        self.load_state_dict(torch.load(pt_fn))

    def get_weights(self) -> torch.tensor:
        return self.coefficients[0]     

class NlbfV1(object):
    def __init__(self):
        self.name = 'apps.wfs.nlbf_v1.NlbfV1'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'非线性波束成形应用 v1 001')
        # 设置中文字体
        plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
        if params['run_mode'] == 1: # 生成数据集
            NlbfV1.generate_ds(
                X_fn = './work/wfs/nlbf/v1/a000/X_8.pt',
                y_fn = './work/wfs/nlbf/v1/a000/y_12_30dB_0_nlbf.pt',
                N_base = 8,
                N_aim = 12,
                ssl_db = 30,
                theta = 0.0
            )
        elif params['run_mode'] == 2: # 模型训练
            NG.rank = 3 # 设置为3阶2次多项式
            NlbfV1.train_main(
                pt_fn='./work/wfs/nlbf/v1/a000/nlbf_v1.pt', 
                X_fn='./work/wfs/nlbf/v1/a000/X_8.pt', 
                y_fn='./work/wfs/nlbf/v1/a000/y_12_30db_0_nlbf.pt'
            )
        elif params['run_mode'] == 3: # 模型训练效果可视化
            NG.rank = 3
            NlbfV1.evaluate(
                pt_fn='./work/wfs/nlbf/v1/a000/nlbf_v1.pt', 
                X_fn='./work/wfs/nlbf/v1/a000/X_8.pt',  
                y_fn='./work/wfs/nlbf/v1/a000/y_12_30db_0_nlbf.pt'
            )
        else:
            print(f'未知运行模式')
    
    @staticmethod
    def train_main(pt_fn:str = './work/wfs/a000/nlbf_v1.pt', X_fn:str = './work/wfs/a000/X_nlbf.pt', y_fn:str = './work/wfs/a000/y_nlbf.pt') -> None:
        '''
        不能收敛的训练过程
        '''
        print(f'{NG.rank}阶多项式回归 v0.0.1')
        # 载入波束成形相关内容和训练目标
        X = torch.load(X_fn, weights_only=True) # 导向向量
        X = X.T
        X = X.to(NG.device)
        y = torch.load(y_fn, weights_only=True).unsqueeze(1)
        y = y.to(NG.device)
        # 创建模型实例
        print(f'### X: {X.shape}; y: {y.shape};')
        model = PolynomialModel1(avec_dim=1, rank=NG.rank, avec=X)
        model = model.to(NG.device)
        # 定义损失函数
        criterion = nn.MSELoss()
        # criterion = nn.L1Loss()
        # 定义优化器
        optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
        # 预热阶段
        print(f'##### 开始预热训练')
        warmup_opt = torch.optim.AdamW(model.parameters(), lr=1e-5, betas=(0.9, 0.99), weight_decay=1e-2)
        warmup_scheduler = ReduceLROnPlateau(warmup_opt, 'min', threshold=1e-4, eps=1e-8)
        warmup_epochs = 5
        for epoch in range(warmup_epochs):
            print(f'    warmup {epoch}......')
            NlbfV1.train_epoch(epoch, model, criterion, warmup_opt, warmup_scheduler, X, y)
        # 正式训练过程
        print(f'##### 开始正式训练')
        optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, betas=(0.9, 0.99), weight_decay=1e-2)
        scheduler = ReduceLROnPlateau(optimizer, 'min', threshold=1e-4, eps=1e-8)
        epochs = 2000000
        best_loss = sys.float_info.max
        improve_threshold = 1e-7
        cumulative_steps = 0
        max_unimproves = 500
        for epoch in range(epochs):
            loss = NlbfV1.train_epoch(epoch, model, criterion, optimizer, scheduler, X, y)
            print(f'epoch_{epoch}: loss={loss};')
            cumulative_steps += 1
            if best_loss > loss:
                if best_loss - loss > improve_threshold:
                    print(f'Save best weights')
                    model.save(pt_fn=pt_fn)
                    cumulative_steps = 0
                    best_loss = loss
            if cumulative_steps > max_unimproves:
                print(f'Earlly Stopping!!!!!!')
                break

    @staticmethod
    def train_epoch(epoch, model, criterion, optimizer, scheduler, X, y_gt):
        batch_size = 10
        max_loss = torch.tensor([0.0]).to(X.device)
        total_loss = torch.tensor([0.0]).to(X.device)
        # 训练过程
        for bidx in range(batch_size):
            loss = None
            y_hat = None
            y_hat = model(X)
            loss = criterion(y_hat, y_gt)
            if loss > max_loss:
                max_loss = loss
            total_loss += loss
            # 清空梯度
            optimizer.zero_grad()
            # 反向传播
            loss.backward()
            # 更新参数
            optimizer.step()
        if scheduler is not None:
            scheduler.step(total_loss)
        return max_loss
            

    @staticmethod
    def evaluate(pt_fn:str = './work/wfs/nlbf_v1.pt', X_fn:str = './work/wfs/X_nlbf.pt', y_fn:str = './work/wfs/y_nlbf.pt') -> None:
        print(f'非线性波束形成简单评估({NG.rank}阶多项式回归) v0.0.2')
        NG.N = 8
        # Tschebyscheff波束形成
        theta = 0 # 扫描角度
        theta_rad = theta / 180 * np.pi
        n_arr = np.arange(NG.N).reshape((-1, 1))
        a = np.exp(1j * 2 * np.pi * np.sin(NG.thetas_rad) * n_arr * NG.d[0] / NG.lambda_)
        w = np.exp(1j * 2 * np.pi * np.sin(theta_rad) * n_arr * NG.d[0] / NG.lambda_)
        B = np.dot(a.transpose(), np.conj(w))
        B = np.abs(B) / np.max(np.abs(B))
        # 载入波束成形相关内容和训练目标
        X = torch.load(X_fn, weights_only=True) # 导向向量
        X = X.T
        X = X.to(NG.device)
        y = torch.load(y_fn, weights_only=True).unsqueeze(1)
        y = y.to(NG.device)
        print(f'### X: {X.shape}, {X.dtype}; ?????????????')
        print(f'### y: {y.shape}, {y.dtype}; {y[0]};')
        # 创建模型实例
        model = PolynomialModel1(avec_dim=1, rank=NG.rank, avec=X)
        model.load(pt_fn = pt_fn) # 载入权重文件
        model = model.to(NG.device)
        y_hat = model(X)
        # 进行可视化
        # 训练后的结果评估
        a_np = X.cpu().numpy()
        # 天线图响应
        y_np = y.detach().cpu().numpy()
        y_hat_np = y_hat.detach().cpu().numpy()
        y_hat_np = np.abs(y_hat_np) / np.max(np.abs(y_hat_np))
        # 绘制波束形成图
        fig, ax = plt.subplots(1, 1, sharex=True, sharey=True, figsize=(7, 2), dpi=300)
        plt.subplots_adjust(wspace=0.1, hspace=0.05)
        ax.plot(NG.thetas[0], 20 * np.log10(y_np), '--', label='12天线')
        ax.plot(NG.thetas[0], 20 * np.log10(y_hat_np), '', label='3阶多项式')
        ax.plot(NG.thetas[0], 20 * np.log10(B), '--', label='8天线')
        ax.legend()
        ax.set_title('Beamforming Result(r=10)', fontsize=12, color='blue', loc='center')
        ax.set_ylim(-60, 1)
        plt.show()

    @staticmethod
    def generate_ds(X_fn:str = './work/wfs/a000/X_nlbf.pt', y_fn:str = './work/wfs/a000/y_nlbf.pt', N_base:int = 8, N_aim:int = 32, ssl_db:float = -30, theta:float = 0.0) -> None:
        '''
        N_base: 基础天线数量
        N_target: 优化目标的天线数量
        ssl_db: 旁瓣水平
        '''
        theta_rad = theta / 180.0 * np.pi
        af = SSW.chebwin(N_aim, at=ssl_db, sym=True)
        phases = np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin(theta_rad))
        w_vec_aim = af * phases # 权重向量
        # 导向向量
        a_vec_aim = np.array([np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_aim) * NG.d[0] * np.sin((ti))) for ti in NG.thetas_rad[0]]).T
        # 天线图响应
        B = np.dot(a_vec_aim.transpose(), np.conj(w_vec_aim))
        B = np.abs(B) / np.max(np.abs(B))
        print(f'w_vec: {w_vec_aim.shape}; a_vec: {a_vec_aim.shape}; B: {B.shape}')
        a_vec_base = np.array([np.exp(1j * 2 * np.pi / NG.lambda_ * np.arange(N_base) * NG.d[0] * np.sin((ti))) for ti in NG.thetas_rad[0]]).T
        X = torch.from_numpy(a_vec_base)
        torch.save(X, X_fn)
        y = torch.from_numpy(B)
        torch.save(y, y_fn)

def main(params:Dict = {}) -> None:
    NlbfV1.startup(params=params)

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--run_mode', action='store',
        type=int, default=1, dest='run_mode',
        help='run mode'
    )
    return parser.parse_args()

if '__main__' == __name__:
    args = parse_args()
    params = vars(args)
    main(params=params)